import numpy as np
import pandas as pd


def sep(label=''):
    print('-' * 32, label, '-' * 32, sep='')


# 1．按照要求，完成朴素贝叶斯底层处理（25分）
# （1）数据预处理
# ①读取bayes_lihang.txt文件（1分）
sep(' ①读取bayes_lihang.txt文件')
df = pd.read_csv(r'../../../../ML_2/bayes/follow_teacher/bayes_lihang.txt')
print(df[:5])

sep('概率存储字典')
# ②创建概率存储字典（1分）
# ③获取所有标签名和特征名称（1分）
# （2）概率计算
# ①计算先验概率的结果并存储（4分）
# ②遍历每个特征，计算每个特征的条件概率（4分）
# ③将上面计算的概率存储的字典中（4分）
lam = 1.0
P_xy = {}
m = len(df)
x = df.iloc[:, :-1]
y = df.iloc[:, -1]
x_names = x.columns
y_cnt_series = y.value_counts()
y_v_len = len(y_cnt_series)
P_y_series = (y_cnt_series + lam) / (m + y_v_len * lam)
for y_v in y_cnt_series.index:
    m_y_v = y_cnt_series[y_v]
    df2 = df.loc[y == y_v]
    for x_name in x_names:
        x_cnt_series = df2[x_name].value_counts()
        x_v_len = len(x_cnt_series)
        k = (y_v, x_name, None)
        P_xy[k] = lam / (m_y_v + x_v_len * lam)
        for x_v in x_cnt_series.index:
            k = (y_v, x_name, x_v)
            P_xy[k] = (x_cnt_series[x_v] + lam) / (m_y_v + x_v_len * lam)


def get_P_of_key(k):
    y_v, x_name, x_v = k
    if k in P_xy:
        return P_xy[k]
    else:
        return P_xy[(y_v, x_name, None)]


# （3）模型预测
# ①打印所有概率的结果（4分）
sep('①打印所有概率的结果')
print(P_y_series)
print(P_xy)


# ②数据测试数据2, s，预测两个类别的概率分别是多少（3分）
sep('②数据测试数据2, s，预测两个类别的概率分别是多少')
x = [2, 'S']
p_arr = []
for y_v in P_y_series.index:
    p = P_y_series[y_v]
    for i, x_v in enumerate(x):
        x_name = x_names[i]
        k = (y_v, x_name, x_v)
        p *= get_P_of_key(k)
    print(f'为类别{y_v}的概率为\t{p}')
    p_arr.append(p)

# ③打印最终的分类类别（3分）
sep('③打印最终的分类类别')
p_arr = np.array(p_arr)
most_prob_idx = p_arr.argsort()[-1]
print(f'最终的分类类别:{P_y_series.index[most_prob_idx]}')
